/*
 * bst/rb_tree.hh - Red-black tree implementation.
 * Copyright (C) 2008  Mikhail Vorozhtsov
 * See the LICENSE section of the README file for details.
 */

/* $Id$ */

#ifndef HEADER_LIBBST_MJ_RB_TREE_HH
#define HEADER_LIBBST_MJ_RB_TREE_HH

#include HEADER (LIBBST_PREFIX,base.hh)

START_NAMESPACE (LIBBST_NS)

template <typename P>
class basic_rb_node_t : public P
{
public:
  INLINE
  basic_rb_node_t *
  get_left () const
  {
    return (basic_rb_node_t *) (this->left.raw & P::BASE_MASK);
  }

  INLINE
  void
  set_left (basic_rb_node_t *x)
  {
    this->left.base = (u64_t) x >> P::BASE_SHIFT;
  }

  INLINE
  bool
  red_p () const
  {
    return (bool) this->left.bit1;
  };

  INLINE
  void
  set_red (bool x)
  {
    this->left.bit1 = x;
  };
};

template <typename NODE, typename TREE, typename GLUE>
struct fwd_rb_tree_t : public base_tree_t<NODE, TREE, GLUE>
{
public:
  template <typename STACK>
  bool
  insert (typename NODE::container_t *x, STACK &stack)
  {
    NODE *parent, *prev_or_next;
    bool left_p;

    if (full_lookup (x, parent, left_p, prev_or_next, stack))
      return false;

    NODE *n = (NODE *) NODE::from_container (x);
    before_insert (n, parent, left_p);
    n->insert (this, parent, left_p, prev_or_next);
    n->set_red (true);

    if (parent)
      stack.pop ();

    while (n != this->get_root ())
      {
        if (parent == NULL)
          parent = stack.pop (this, n);

        if (!parent->red_p ())
          return true;

        NODE *grandparent = stack.pop (this, parent);

        if (parent == grandparent->get_left ())
          {
            NODE *uncle = grandparent->get_right ();

            if (uncle && uncle->red_p ())
              {
                parent->set_red (false);
                uncle->set_red (false);
                grandparent->set_red (true);
                n = grandparent;
                parent = NULL;
                continue;
              }

            if (n == parent->get_right ())
              {
                parent->rotate_left (this, grandparent);
                NODE *tmp = n;
                n = parent;
                parent = tmp;
              }
        
            parent->set_red (false);
            grandparent->set_red (true);
            grandparent->rotate_right (this,
              grandparent == this->get_root () ? NULL
              : stack.pop (this, grandparent));
            return true;
          }
        else
          {
            NODE *uncle = grandparent->get_left ();

            if (uncle && uncle->red_p ())
              {
                parent->set_red (false);
                uncle->set_red (false);
                grandparent->set_red (true);
                n = grandparent;
                parent = NULL;
                continue;
              }

            if (n == parent->get_left ())
              {
                parent->rotate_right (this, grandparent);
                NODE *tmp = n;
                n = parent;
                parent = tmp;
              }

            parent->set_red (false);
            grandparent->set_red (true);
            grandparent->rotate_left (this,
              grandparent == this->get_root () ? NULL
              : stack.pop (this, grandparent));
            return true;
          }
      }

    n->set_red (false);
    return true;
  }

  template <typename KEY, typename CMP, typename STACK>
  typename NODE::container_t *
  remove_key (const KEY &key, STACK &stack)
  {
    NODE *parent, *prev_or_next;
    bool left_p;
    typename NODE::container_t *x
      = this->template full_lookup_key<KEY, CMP, STACK>
                         (key, parent, left_p, prev_or_next, stack);

    if (!x)
      return NULL;

    NODE *n = (NODE *) NODE::from_container (x);
    NODE *left = (NODE *) n->get_left ();
    NODE *right = (NODE *) n->get_right ();
    NODE *child;
    NODE *grandparent = NULL;

    if (left)
      {
        NODE *pred_parent = n;
        NODE *pred = left->get_max_and_replace (pred_parent, stack);
        NODE *succ = NULL;
        child = n->replace_with_pred (this, parent, left_p, pred, pred_parent,
                                      succ);

        pred->set_red (n->red_p ());
        after_remove (n, parent, pred, right != NULL, succ);

        if (pred_parent == n)
          {
            grandparent = parent;
            parent = pred;

            if (grandparent)
              stack.pop ();
          }
        else
          parent = pred_parent;
      }
    else if (right)
      {
        NODE *succ_parent = n;
        NODE *succ = right->get_min_and_replace (succ_parent, stack);
        child = n->replace_with_succ (this, parent, left_p, succ, succ_parent,
                                      prev_or_next);

        succ->set_red (n->red_p ());
        after_remove (n, parent, NULL, true, succ);

        if (succ_parent == n)
          {
            grandparent = parent;
            parent = succ;

            if (grandparent)
              stack.pop ();
          }
        else
          parent = succ_parent;
      }
    else
      {
        n->remove (this, parent, left_p, prev_or_next);
        after_remove (n, parent, NULL, false, NULL);
        child = NULL;
      }

    if (n->red_p ())
      return x;

    if (parent)
      stack.pop ();

    for (n = child; n != this->get_root () && (!n || !n->red_p ());)
      {
        NODE *s, *next_parent;

        if (grandparent == NULL)
          grandparent = parent == this->get_root () ? NULL
                        : stack.pop (this, parent);

        if (n == parent->get_left ())
          {
            s = (NODE *) parent->get_right ();

            if (s == NULL)
              {
                n = parent;
                parent = grandparent;
                grandparent = NULL;
                continue;
              }

            if (s->red_p ())
              {
                s->set_red (false);
                parent->set_red (true);
                parent->rotate_left (this, grandparent);
                next_parent = s;
                s = (NODE *) parent->get_right ();

                if (s == NULL)
                  {
                    n = parent;
                    parent = next_parent;
                    continue;
                  }
              }
            else
              {
                next_parent = grandparent;
                grandparent = NULL;
              }

            NODE *left = (NODE *) s->get_left ();
            NODE *right = (NODE *) s->get_right ();

            if (!right || !right->red_p ())
              {
                if (!left || !left->red_p ())
                  {
                    s->set_red (true);
                    n = parent;
                    parent = next_parent;
                    continue;
                  }

                if (left)
                  left->set_red (false);

                s->set_red (true);
                s->rotate_right (this, parent);
                right = s;
                s = left;
              }

            s->set_red (parent->red_p ());
            parent->set_red (false);

            if (right)
              right->set_red (false);

            parent->rotate_left (this, next_parent);
            n = (NODE *) this->get_root ();
            break;
          }
        else
          {
            s = (NODE *) parent->get_left ();

            if (s == NULL)
              {
                n = parent;
                parent = grandparent;
                grandparent = NULL;
                continue;
              }

            if (s->red_p ())
              {
                s->set_red (false);
                parent->set_red (true);
                parent->rotate_right (this, grandparent);
                next_parent = s;
                s = (NODE *) parent->get_left ();

                if (s == NULL)
                  {
                    n = parent;
                    parent = next_parent;
                    continue;
                  }
              }
            else
              {
                next_parent = grandparent;
                grandparent = NULL;
              }

            NODE *left = (NODE *) s->get_left ();
            NODE *right = (NODE *) s->get_right ();

            if (!left || !left->red_p ())
              {
                if (!right || !right->red_p ())
                  {
                    s->set_red (true);
                    n = parent;
                    parent = next_parent;
                    continue;
                  }

                if (right)
                  right->set_red (false);

                s->set_red (true);
                s->rotate_left (this, parent);
                left = s;
                s = right;
              }

            s->set_red (parent->red_p ());
            parent->set_red (false);

            if (left)
              left->set_red (false);

            parent->rotate_right (this, next_parent);
            n = (NODE *) this->get_root ();
            break;
          }

      }

    if (n)
      n->set_red (false);

    return x;
  }

public:
  void
  replace (typename NODE::container_t *x, typename NODE::container_t *with)
  {
    NODE *node = NODE::from_container (x);
    NODE *with_node = NODE::from_container (with);

    with_node->set_red (node->red_p ());
    base_tree_t<NODE, TREE, GLUE>::replace (x, with);
  }
};

template <typename NODE, typename TREE = basic_tree_t, typename GLUE = glue_t>
struct rb_tree_t
  : public wrap_tree_t<ord_lookup_tree_t<fwd_rb_tree_t<NODE, TREE, GLUE> > >
{
};

END_NAMESPACE

#endif /* HEADER_LIBBST_MJ_RB_TREE_HH */

